#include "ompl/base/goals/GoalSampleableRegion.h"
#include "LazyRRTIterationCount.h"

using namespace ompl;

LazyRRTIterationCount::LazyRRTIterationCount(const ompl::base::SpaceInformationPtr &si)
: ompl::geometric::LazyRRT(si)
{
}

LazyRRTIterationCount::~LazyRRTIterationCount(void)
{
}

void LazyRRTIterationCount::clear(void)
{
	ompl::geometric::LazyRRT::clear();
	uiNumIterations_ = 0;
}

// ++uiNumIterations_ is added
ompl::base::PlannerStatus LazyRRTIterationCount::solve(const ompl::base::PlannerTerminationCondition &ptc)
{
	checkValidity();
    base::Goal                 *goal   = pdef_->getGoal().get();
    base::GoalSampleableRegion *goal_s = dynamic_cast<base::GoalSampleableRegion*>(goal);

    while (const base::State *st = pis_.nextStart())
    {
        Motion *motion = new Motion(si_);
        si_->copyState(motion->state, st);
        motion->valid = true;
        nn_->add(motion);
    }

    if (nn_->size() == 0)
    {
        OMPL_ERROR("There are no valid initial states!");
        return base::PlannerStatus::INVALID_START;
    }

    if (!sampler_)
        sampler_ = si_->allocStateSampler();

    OMPL_INFORM("Starting with %u states", nn_->size());

    Motion *solution = NULL;
    double  distsol  = -1.0;
    Motion *rmotion  = new Motion(si_);
    base::State *rstate = rmotion->state;
    base::State *xstate = si_->allocState();

    bool solutionFound = false;

    while (ptc == false && !solutionFound)
    {
		++uiNumIterations_;

        /* sample random state (with goal biasing) */
        if (goal_s && rng_.uniform01() < goalBias_ && goal_s->canSample())
            goal_s->sampleGoal(rstate);
        else
            sampler_->sampleUniform(rstate);

        /* find closest state in the tree */
        Motion *nmotion = nn_->nearest(rmotion);
        assert(nmotion != rmotion);
        base::State *dstate = rstate;

        /* find state to add */
        double d = si_->distance(nmotion->state, rstate);
        if (d > maxDistance_)
        {
            si_->getStateSpace()->interpolate(nmotion->state, rstate, maxDistance_ / d, xstate);
            dstate = xstate;
        }

        /* create a motion */
        Motion *motion = new Motion(si_);
        si_->copyState(motion->state, dstate);
        motion->parent = nmotion;
        nmotion->children.push_back(motion);
        nn_->add(motion);

        double dist = 0.0;
        if (goal->isSatisfied(motion->state, &dist))
        {
            distsol = dist;
            solution = motion;
            solutionFound = true;
            lastGoalMotion_ = solution;

            // Check that the solution is valid:
            // construct the solution path
            std::vector<Motion*> mpath;
            while (solution != NULL)
            {
                mpath.push_back(solution);
                solution = solution->parent;
            }

            // check each segment along the path for validity
            for (int i = mpath.size() - 1 ; i >= 0 && solutionFound; --i)
                if (!mpath[i]->valid)
                {
                    if (si_->checkMotion(mpath[i]->parent->state, mpath[i]->state))
                        mpath[i]->valid = true;
                    else
                    {
                        removeMotion(mpath[i]);
                        solutionFound = false;
                        lastGoalMotion_ = NULL;
                    }
                }

            if (solutionFound)
            {
                // set the solution path
                ompl::geometric::PathGeometric *path = new ompl::geometric::PathGeometric(si_);
                for (int i = mpath.size() - 1 ; i >= 0 ; --i)
                    path->append(mpath[i]->state);

                pdef_->addSolutionPath(base::PathPtr(path), false, distsol);
            }
        }
    }

    si_->freeState(xstate);
    si_->freeState(rstate);
    delete rmotion;

    OMPL_INFORM("Created %u states", nn_->size());

    return solutionFound ?  base::PlannerStatus::EXACT_SOLUTION : base::PlannerStatus::TIMEOUT;
}